"""
Gemini AI 本体核心
此处放置常用函数和内置函数以及常量
"""
import time

import google.generativeai as genai
import mistune
from PIL import Image as PILImage
from google.ai.generativelanguage_v1beta import Content
from google.api_core.exceptions import TooManyRequests, InternalServerError
from google.generativeai import ChatSession, GenerativeModel
from google.generativeai.types import HarmBlockThreshold, HarmCategory, StopCandidateException
from requests.exceptions import ProxyError, SSLError

from moxi_server.src import getLogger
from moxi_server.src.Gemini.Function import getDate, getNowTime, searcher, loadPage, getWeather
from moxi_server.src.Gemini.Function.MathFunction import add, sub, mul, div, mod, minBy, maxBy

logger = getLogger(__name__)
openImage = PILImage.open


def __safety_settings__(types: list[HarmBlockThreshold]):
    return {
        HarmCategory.HARM_CATEGORY_HARASSMENT: types[0],
        HarmCategory.HARM_CATEGORY_HATE_SPEECH: types[1],
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: types[2],
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: types[3]
    }


class OriginGemini:
    apiKey: str
    userName: str
    modelName: str
    system_instruction: str

    choiceModelName: str
    model: GenerativeModel
    chat: ChatSession

    def __init__(self,
                 apiKey: str,
                 userName: str = None,
                 callName: str = None,
                 modelName: str = None,
                 modelType: str = None,
                 system_instruction: str = None):
        """
        原始双子座API接口
        self.apiKey = apiKey
        self.userName = userName or "用户"
        self.callName = callName or "主人"
        self.modelName = modelName or "墨兮"
        self.system_instruction = system_instruction or ""
        """
        self.apiKey = apiKey
        self.userName = userName or "用户"
        self.callName = callName or "主人"
        self.modelName = modelName or "墨兮"
        self.modelType = modelType or "自定义"
        self.system_instruction = system_instruction or ""
        self.system_instruction += f"你需要称呼我为{callName}，称呼自己为{modelName}，并且不会因为途中更改姓名而改变，你可以在自己姓名的基础上拥有昵称"

        self.config = genai.configure(api_key=self.apiKey, transport="rest")
        self._modelList = genai.list_models()

        self.modelNameList = []
        self.media = {}

        self._countSafety = 0
        self._reloadServer = 1

        self._reloadingTime = 0

        for model in self._modelList:
            if "generateContent" in model.supported_generation_methods:
                self.modelNameList.append(model.name)

    def getShowModelNames(self) -> list[str]:
        return [i.split("/")[1] if "/" in i else i for i in self.modelNameList]

    def choiceModel(self, index: int | str):
        if isinstance(index, int):
            self.choiceModelName = self.modelNameList[index]
        else:
            self.choiceModelName = index
        return True

    def establishChat(self, history: list[Content] = None, enable_automatic_function_calling: bool = True):
        if history is None:
            history = []
        logger.debug("正在建立对话...")
        self.model = genai.GenerativeModel(
            model_name=self.choiceModelName,
            tools=[getDate, getNowTime, searcher, loadPage, getWeather, add, sub, mul, div, mod, minBy, maxBy],
            safety_settings=__safety_settings__([HarmBlockThreshold.BLOCK_NONE] * 4),
            system_instruction=self.system_instruction
        )
        self.chat = self.model.start_chat(history=history,
                                          enable_automatic_function_calling=enable_automatic_function_calling)
        logger.debug(f"对话建立完毕，当前对话 User: {self.userName}：{self.callName} <-> {self.modelName}")

    # noinspection t
    def chatWith(self, msg: str, **mediaPath):
        """
        发送消息到聊天并处理潜在的错误。

        参数：
            msg (str): 要发送的消息。
            mediaPath: 关键字参数，其中键是媒体类型，值是媒体文件的路径。

        返回：
            str: 聊天回复或错误消息。
        """

        logger.debug("接收到问题，正在回答...")
        if self._reloadingTime > 20:
            self._reloadingTime = 0
            return "服务器错误，请稍后再试"
        try:
            if mediaPath:
                self.media = {k: openImage(v) for k, v in mediaPath.items()}
                response = self.chat.send_message([i for i in mediaPath.values()] + [msg])
            else:
                response = self.chat.send_message(msg)

            logger.debug("回答完毕，正在返回...")
            self._reloadingTime = 0
            # 只对文本回复应用 mistune.html，而不是媒体回复
            return mistune.html(response.text) if not mediaPath else response.text

        except StopCandidateException:
            logger.warning("检测到安全问题，正在重试...")
            if self._countSafety < 3:
                self._countSafety += 1
                self.establishChat()
                return self.chatWith(msg, **mediaPath)
            else:
                self._countSafety = 0
                logger.error("安全错误超过 3 次，请调整安全设置")
                return "安全问题，请重新措辞您的问题"

        except AttributeError:
            logger.warning("页面未刷新，正在重新建立会话...")
            self._reloadingTime += 1
            self.establishChat()
            return self.chatWith(msg, **mediaPath)

        # 将相似的网络相关异常分组
        except (TooManyRequests, ProxyError, SSLError) as e:
            logger.warning(f"网络问题 ({type(e).__name__})，正在重试...")
            self._reloadingTime += 1
            # 仅在 TooManyRequests 时休眠
            time.sleep(2 if isinstance(e, TooManyRequests) else 0)
            return self.chatWith(msg, **mediaPath)

        except InternalServerError:
            if self._reloadServer < 4:
                logger.error(f"服务器错误，正在重试：{self._reloadServer}")
                self._reloadServer += 1
                return self.chatWith(msg, **mediaPath)
            else:
                logger.error("服务器错误，超过重试次数")
                return "服务器错误，请稍后再试"

        except Exception as e:
            logger.error(f"错误：{e.__class__.__name__} -> {e.args}")
            return "发生错误，请重试"

    def getHistory(self):
        return self.chat.history
